# -*- coding: utf-8 -*-
# +
# -*- coding: utf-8 -*-
import os
import xgboost as xgb
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

from tools.config import CONFIG
from tools import functions
# from tools.data_processing import *

class PyModel:
    def __init__(self, station, params={}):
        self.params = CONFIG().PARAMS_DICT[station]
        if params: self.params.update(params)
        self.model = {}
        self.station = station
        self.basePath = os.path.split(os.path.realpath(__file__))[0]
        self.save_path = os.path.join(self.basePath, 'save_models')
        self.pic_save_path = os.path.join(self.basePath, 'figures')
        functions.ensure_folder(self.save_path)
        functions.ensure_folder(self.pic_save_path)
        self.labels = CONFIG().LABELS

    def fit(self, df_X, df_y, show_our_score=True):
        
        ret = pd.DataFrame()
        for y_label in self.labels:
            print('[INFO] training:', self.station, y_label)
            X_train, X_test, y_train, y_test = train_test_split(df_X, df_y[y_label], test_size=0.1, shuffle=False, random_state=1234565)
            self.model[y_label] = xgb.XGBRegressor(**self.params)
            self.model[y_label].fit(X_train, y_train, eval_set=[(X_train,y_train),(X_test, y_test)], eval_metric="rmse", verbose=0)
            results = self.model[y_label].evals_result()
            save_path_loss = self.pic_save_path+'/'+ self.station + '_' + y_label + '_loss.png'
            functions.show_figure(results['validation_0']['rmse'], 
                                  results['validation_1']['rmse'], 
                                  title=self.station+'-'+y_label+'-'+'loss', 
                                  save_path=save_path_loss,
                                  labels=['train','test'],
                                  xlabel='step',ylabel='loss')
            self.save_(y_label,os.path.join(self.save_path, '%s_%s.model'%(self.station, y_label)))
            if show_our_score:
                # 测试误差
                save_path_test = os.path.join(self.pic_save_path, self.station + '_' + y_label + '_fitness_test.png')
                y_pred = self.model[y_label].predict(X_test)
                functions.show_figure(y_test, y_pred, 
                                      title=self.station+'-'+y_label+'valid-fitness', 
                                      save_path=save_path_test)
                scores = {'场站':[self.station], 'label':y_label,}
                scores.update(functions.my_score(y_pred, y_test, 'test_'))
                # 训练误差
                save_path_train = os.path.join(self.pic_save_path, self.station + '_' + y_label + '_fitness_train.png')
                y_pred = self.model[y_label].predict(X_train)
                functions.show_figure(y_train, y_pred, 
                                      title=self.station+':'+y_label+'train-fitness', 
                                      save_path=save_path_train)
                functions.concat_figures([save_path_loss,save_path_train,save_path_test],save_path=os.path.join(self.basePath, 'concated_figures', self.station + '_' + y_label+'.png'))
                scores.update(functions.my_score(y_pred, y_train, 'train_'))
                ret = pd.concat([ret, pd.DataFrame(scores)])
    #             print(self.station, scores)
        return ret

    def predict(self, df_X):
        ret = pd.DataFrame()
        if not self.model:
            self.load()
#         df = feature_encode(df, match_type, y_label=None, station=self.station) #只剩使用的特征和标签（标签原封不动）
        for y_label in self.labels:
            y_pred = pd.DataFrame(self.model[y_label].predict(df_X), columns=[y_label])
            ret = pd.concat([ret, y_pred], axis=1)
        return ret

    def save_(self, y_label, model_dir):
        self.model[y_label].save_model(model_dir)

    def load(self):
        for y_label in self.labels:
            self.model[y_label] = xgb.XGBRegressor()
            self.model[y_label].load_model(os.path.join(self.save_path, '%s_%s.model'%(self.station, y_label)))

# -


